{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8f360643-0fc7-4131-8070-c21743d50fb0",
   "metadata": {},
   "source": [
    "# Test the Accelerators on Simple Examples\n",
    "\n",
    "This notebook tests each of the four accelerators described in Sections 3, 4, and 5 and evaluated in Section 7 on a small example. This ensures that all compilation and modeling are performed correctly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a112574-bdca-4fd0-813a-0b2e05db65d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "sys.path.insert(0, \"/home/workspace\")\n",
    "\n",
    "from fibertree import Tensor\n",
    "\n",
    "import scripts.extensor as extensor\n",
    "import scripts.extensor_energy as extensor_energy\n",
    "import scripts.gamma as gamma\n",
    "import scripts.outerspace as outerspace\n",
    "import scripts.sigma as sigma"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb554a86-2f27-414b-9a10-0fff678b8550",
   "metadata": {},
   "outputs": [],
   "source": [
    "K = 8\n",
    "M = 6\n",
    "N = 7\n",
    "density = 0.5\n",
    "\n",
    "# Create the tensors\n",
    "A_KM = Tensor.fromRandom(rank_ids=[\"K\", \"M\"], shape=[K, M], density=[0.9, density], seed=0)\n",
    "B_KN = Tensor.fromRandom(rank_ids=[\"K\", \"N\"], shape=[K, N], density=[0.9, density], seed=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65540a7c-5ec4-460f-ba28-ace377930736",
   "metadata": {},
   "outputs": [],
   "source": [
    "extensor.check(extensor.run(A_KM, B_KN, 2, 2, 2, 2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96f4d067-f206-4835-b8c5-da16c5abc645",
   "metadata": {},
   "outputs": [],
   "source": [
    "extensor_energy.check(extensor_energy.run(A_KM, B_KN, 2, 2, 2, 2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34b8f9b2-419f-4a02-87ca-0066943ae4c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "A_MK = A_KM.swizzleRanks([\"M\", \"K\"])\n",
    "gamma.check(gamma.run(A_MK, B_KN))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe9a6dad-95ac-40d3-8164-93a4de12ca77",
   "metadata": {},
   "outputs": [],
   "source": [
    "outerspace.check(outerspace.run(A_KM, B_KN))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0930b023-2933-45e5-ae20-96d342b22e58",
   "metadata": {},
   "outputs": [],
   "source": [
    "sigma.check(sigma.run(A_KM, B_KN))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "affc486b-a9a2-4c7b-9cfc-1f49fa49e92e",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
